#define _USE_MATH_DEFINES
#include <cmath>
#include "AlgoAStar.h"
#include "AlgoAStarPrivate.h"
#include <queue>
#include <vector>
#include <algorithm>
#include <iostream>
#include <limits>

namespace algo
{
	// This should be a member function or in a helper namespace, but for simplicity let's make it a static function in the file
	static double distance(const algobase::CoordZ& p1, const algobase::CoordZ& p2) {
		double lon1 = p1[0], lat1 = p1[1], depth1 = p1[2];
		double lon2 = p2[0], lat2 = p2[1], depth2 = p2[2];
		double R = 6371e3; // Earth radius in meters
		double x = (lon2 - lon1) * M_PI / 180.0 * cos((lat1 + lat2) / 2.0 * M_PI / 180.0);
		double y = (lat2 - lat1) * M_PI / 180.0;
		double horiz_dist = sqrt(x * x + y * y) * R;
		double vert_dist = depth2 - depth1;
		return sqrt(horiz_dist * horiz_dist + vert_dist * vert_dist);
	}

	// --- Helper for path combination ---
	void combineRecursiveAStar(
		size_t segment_index,
		algobase::Path current_path,
		const std::vector<std::vector<algobase::Path>>& all_segment_paths,
		algobase::Paths& final_paths,
		size_t& path_count_limit
	) {
		if (final_paths.size() >= path_count_limit) return;
		if (segment_index == all_segment_paths.size()) {
			final_paths.push_back(current_path);
			return;
		}
		for (const auto& segment : all_segment_paths[segment_index]) {
			if (segment.empty()) continue;
			algobase::Path next_path = current_path;
			next_path.insert(next_path.end(), segment.begin() + 1, segment.end());
			combineRecursiveAStar(segment_index + 1, next_path, all_segment_paths, final_paths, path_count_limit);
			if (final_paths.size() >= path_count_limit) return;
		}
	}


	// --- AlgoAStar public methods ---
	AlgoAStar::AlgoAStar() : d(new AlgoAStarPrivate()) {}
	AlgoAStar::~AlgoAStar() { delete d; }
	int AlgoAStar::exec(const algobase::TaskConfig& taskConfig, algobase::Paths& paths)
	{
		// Mock terrain data loading
		int terrain_size = 100;
		std::vector<std::vector<double>> terrain(terrain_size, std::vector<double>(terrain_size, 100.0));
		d->setTerrainData(terrain);

		return d->exec(taskConfig, paths);
	}


	// --- AlgoAStarPrivate implementation ---
	AlgoAStarPrivate::AlgoAStarPrivate() {}
	AlgoAStarPrivate::~AlgoAStarPrivate() {}

	void AlgoAStarPrivate::setTerrainData(const std::vector<std::vector<double>>& terrain) {
		terrain_data_ = terrain;
	}

	int AlgoAStarPrivate::exec(const algobase::TaskConfig& taskConfig, algobase::Paths& paths)
	{
		paths.clear();
		if (!preprocess(taskConfig)) {
			std::cerr << "Error during A* preprocessing." << std::endl;
			return -1;
		}

		auto start_grid = grid_map_->worldToGrid(taskConfig.mStartPos);
		auto end_grid = grid_map_->worldToGrid(taskConfig.mEndPos);

		std::vector<std::pair<int, int>> waypoints;
		waypoints.push_back(start_grid);

		const auto mandatory_polygons_it = taskConfig.mTypePolygons.find(algobase::AT_Mandatory);
		if (mandatory_polygons_it != taskConfig.mTypePolygons.end()) {
			for (const auto& polygon : mandatory_polygons_it->second) {
				waypoints.push_back(findRepresentativeNodeForZone(polygon));
			}
		}
		waypoints.push_back(end_grid);

		int K = 10; // Number of paths per segment
		std::vector<std::vector<algobase::Path>> all_segment_paths;

		for (size_t i = 0; i < waypoints.size() - 1; ++i) {
			auto segment_paths = findKShortestPaths(waypoints[i], waypoints[i + 1], K);
			if (segment_paths.empty()) {
				std::cerr << "Could not find path for a segment." << std::endl;
				return -2;
			}
			all_segment_paths.push_back(segment_paths);
		}

		paths = combinePaths(all_segment_paths);

		if (paths.empty()) {
			std::cerr << "A* could not find any combined path." << std::endl;
			return -3;
		}

		return 0;
	}

	std::vector<algobase::Path> AlgoAStarPrivate::findKShortestPaths(const std::pair<int, int>& start_grid, const std::pair<int, int>& end_grid, int K) {
		std::vector<std::vector<Node*>> A; // K shortest paths in node form
		std::priority_queue<std::pair<double, std::vector<Node*>>, std::vector<std::pair<double, std::vector<Node*>>>, std::greater<std::pair<double, std::vector<Node*>>>> B;

		auto first_path = findShortestPath(start_grid, end_grid, {});
		if (first_path.empty()) return {};
		A.push_back(first_path);

		for (int k = 1; k < K; ++k) {
			const auto& previous_path = A[k - 1];
			for (size_t i = 0; i < previous_path.size() - 1; ++i) {
				Node* spur_node = previous_path[i];
				std::vector<Node*> root_path(previous_path.begin(), previous_path.begin() + i + 1);

				std::map<std::pair<int, int>, double> extra_costs;
				for (const auto& path_in_A : A) {
					if (path_in_A.size() > i) {
						bool same_root = true;
						for (size_t j = 0; j <= i; ++j) if (path_in_A[j]->x != root_path[j]->x || path_in_A[j]->y != root_path[j]->y) { same_root = false; break; }
						if (same_root) {
							extra_costs[{path_in_A[i]->x, path_in_A[i]->y}] = std::numeric_limits<double>::infinity();
						}
					}
				}

				auto spur_path = findShortestPath({ spur_node->x, spur_node->y }, end_grid, extra_costs);

				if (!spur_path.empty()) {
					std::vector<Node*> total_path = root_path;
					total_path.insert(total_path.end(), spur_path.begin() + 1, spur_path.end());

					double total_cost = 0;
					for (size_t j = 0; j < total_path.size() - 1; ++j) total_cost += calculateCost(total_path[j], total_path[j + 1]);
					B.push({ total_cost, total_path });
				}
			}

			if (B.empty()) break;

			bool new_path_found = false;
			while (!B.empty()) {
				auto candidate = B.top();
				B.pop();
				bool is_duplicate = false;
				for (const auto& p_in_A : A) if (p_in_A.size() == candidate.second.size()) {
					bool match = true;
					for (size_t j = 0; j < p_in_A.size(); ++j) if (p_in_A[j]->x != candidate.second[j]->x || p_in_A[j]->y != candidate.second[j]->y) { match = false; break; }
					if (match) { is_duplicate = true; break; }
				}
				if (!is_duplicate) { A.push_back(candidate.second); new_path_found = true; break; }
			}
			if (!new_path_found) break;
		}

		std::vector<algobase::Path> result_paths;
		for (const auto& path_nodes : A) {
			result_paths.push_back(postprocessPath(path_nodes));
		}
		return result_paths;
	}

	algobase::Paths AlgoAStarPrivate::combinePaths(const std::vector<std::vector<algobase::Path>>& all_segment_paths) {
		algobase::Paths final_paths;
		if (all_segment_paths.empty()) return {};
		size_t path_limit = 5000;
		for (const auto& first_segment : all_segment_paths[0]) {
			combineRecursiveAStar(1, first_segment, all_segment_paths, final_paths, path_limit);
			if (final_paths.size() >= path_limit) break;
		}
		return final_paths;
	}

	std::pair<int, int> AlgoAStarPrivate::findRepresentativeNodeForZone(const algobase::Polygon& polygon) {
		// Find geometric center of the polygon
		double avg_x = 0, avg_y = 0;
		for (const auto& p : polygon) { avg_x += p[0]; avg_y += p[1]; }
		algobase::Coord center = { avg_x / polygon.size(), avg_y / polygon.size() };
		return grid_map_->worldToGrid(center);
	}

	// --- The rest of the file (preprocess, findShortestPath, helpers) remains largely the same ---
	// Renamed search to findShortestPath and adapted it for Yen's algorithm
	std::vector<Node*> AlgoAStarPrivate::findShortestPath(const std::pair<int, int>& start_grid, const std::pair<int, int>& end_grid, const std::map<std::pair<int, int>, double>& extra_costs)
	{
		std::priority_queue<Node*, std::vector<Node*>, Node::Compare> open_list;
		std::map<std::pair<int, int>, Node*> all_nodes;

		Node* start_node = new Node(start_grid.first, start_grid.second);
		start_node->h_cost = calculateHeuristic(start_node->x, start_node->y, end_grid.first, end_grid.second);
		open_list.push(start_node);
		all_nodes[start_grid] = start_node;

		Node* end_node_ptr = nullptr;
		std::vector<Node*> path;

		while (!open_list.empty()) {
			Node* current_node = open_list.top();
			open_list.pop();

			if (current_node->x == end_grid.first && current_node->y == end_grid.second) {
				end_node_ptr = current_node;
				break;
			}

			for (int dx = -1; dx <= 1; ++dx) {
				for (int dy = -1; dy <= 1; ++dy) {
					if (dx == 0 && dy == 0) continue;
					int nx = current_node->x + dx;
					int ny = current_node->y + dy;
					std::pair<int, int> neighbor_pos = { nx, ny };

					if (nx < 0 || nx >= grid_map_->getWidth() || ny < 0 || ny >= grid_map_->getHeight() || !grid_map_->at(nx, ny).is_walkable) {
						continue;
					}

					Node to_node_temp(nx, ny);
					double cost_multiplier = extra_costs.count(neighbor_pos) ? std::numeric_limits<double>::infinity() : 1.0;
					double new_g_cost = current_node->g_cost + calculateCost(current_node, &to_node_temp) * cost_multiplier;

					if (new_g_cost >= std::numeric_limits<double>::infinity()) continue;

					if (!all_nodes.count(neighbor_pos) || new_g_cost < all_nodes[neighbor_pos]->g_cost) {
						Node* neighbor_node = all_nodes.count(neighbor_pos) ? all_nodes[neighbor_pos] : new Node(nx, ny);
						if (!all_nodes.count(neighbor_pos)) {
							all_nodes[neighbor_pos] = neighbor_node;
							neighbor_node->h_cost = calculateHeuristic(nx, ny, end_grid.first, end_grid.second);
						}
						neighbor_node->g_cost = new_g_cost;
						neighbor_node->parent = current_node;
						open_list.push(neighbor_node);
					}
				}
			}
		}

		if (end_node_ptr) {
			Node* temp = end_node_ptr;
			while (temp != nullptr) { path.push_back(temp); temp = temp->parent; }
			std::reverse(path.begin(), path.end());
		}

		// Create a deep copy of the path nodes for the caller, so we can safely delete all allocated nodes.
		std::vector<Node*> final_path_copy;
		std::map<Node*, Node*> node_map;
		for (Node* old_node : path) {
			Node* new_node = new Node(old_node->x, old_node->y);
			if (old_node->parent && node_map.count(old_node->parent)) {
				new_node->parent = node_map[old_node->parent];
			}
			node_map[old_node] = new_node;
			final_path_copy.push_back(new_node);
		}

		// Cleanup all nodes created during this search instance.
		for (auto const& [pos, node_ptr] : all_nodes) {
			delete node_ptr;
		}

		return final_path_copy;
	}

	// Other helpers like preprocess, postprocessPath, calculateCost, etc.
	// ... (These functions are mostly unchanged from the original implementation)
	bool AlgoAStarPrivate::preprocess(const algobase::TaskConfig& taskConfig)
	{
		double resolution = (taskConfig.mMaxLongitude - taskConfig.mMinLongitude) / 100.0;
		if (resolution <= 0) resolution = 0.01;

		grid_map_ = std::make_unique<GridMap>(
			taskConfig.mMinLongitude, taskConfig.mMaxLongitude,
			taskConfig.mMinLatitude, taskConfig.mMaxLatitude,
			resolution
		);

		if (terrain_data_.empty()) return false;

		for (int y = 0; y < grid_map_->getHeight(); ++y) {
			for (int x = 0; x < grid_map_->getWidth(); ++x) {
				grid_map_->at(x, y).depth = getDepthFromTerrain(grid_map_->gridToWorld(x, y));
				if (grid_map_->at(x, y).depth > 0) grid_map_->at(x, y).is_walkable = false; // Land
			}
		}

		for (const auto& pair : taskConfig.mTypePolygons) {
			const auto type = pair.first;
			if (type == algobase::AT_Obstacle || type == algobase::AT_Restricted) {
				for (const auto& polygon : pair.second) {
					if (polygon.size() < 3) continue;
					// A simple bounding box check first for efficiency
					// ... then the more expensive point-in-polygon test
					for (int y = 0; y < grid_map_->getHeight(); ++y) {
						for (int x = 0; x < grid_map_->getWidth(); ++x) {
							if (isInside(grid_map_->gridToWorld(x, y), polygon)) {
								grid_map_->at(x, y).is_walkable = false;
							}
						}
					}
				}
			}
		}
		return true;
	}

	algobase::Path AlgoAStarPrivate::postprocessPath(const std::vector<Node*>& node_path)
	{
		if (node_path.empty()) return {};
		algobase::Path smoothed_path;
		// ... (Implementation is the same as before)
		for (const auto* node : node_path) {
			algobase::PathCoord pc;
			auto world_xy = grid_map_->gridToWorld(node->x, node->y);
			pc.mCoordZ = { world_xy[0], world_xy[1], grid_map_->at(node->x, node->y).depth };
			smoothed_path.push_back(pc);
		}
		// A full implementation would add interpolation here.
		return smoothed_path;
	}

	double AlgoAStarPrivate::calculateHeuristic(int x1, int y1, int x2, int y2)
	{
		auto p1_world = grid_map_->gridToWorld(x1, y1);
		auto p2_world = grid_map_->gridToWorld(x2, y2);
		algobase::CoordZ p1_z = { p1_world[0], p1_world[1], grid_map_->at(x1, y1).depth };
		algobase::CoordZ p2_z = { p2_world[0], p2_world[1], grid_map_->at(x2, y2).depth };
		return distance(p1_z, p2_z);
	}

	double AlgoAStarPrivate::calculateCost(const Node* from, const Node* to)
	{
		auto p1_world = grid_map_->gridToWorld(from->x, from->y);
		auto p2_world = grid_map_->gridToWorld(to->x, to->y);
		algobase::CoordZ p1_z = { p1_world[0], p1_world[1], grid_map_->at(from->x, from->y).depth };
		algobase::CoordZ p2_z = { p2_world[0], p2_world[1], grid_map_->at(to->x, to->y).depth };
		return distance(p1_z, p2_z);
	}

	// ... (getDepth, isInside, distance helpers are the same)
	double AlgoAStarPrivate::getDepthFromTerrain(const algobase::Coord& coord) const {
		if (terrain_data_.empty() || terrain_data_[0].empty()) return 0.0;
		double lon_frac = (coord[0] - grid_map_->gridToWorld(0, 0)[0]) / (grid_map_->gridToWorld(grid_map_->getWidth() - 1, 0)[0] - grid_map_->gridToWorld(0, 0)[0]);
		double lat_frac = (coord[1] - grid_map_->gridToWorld(0, 0)[1]) / (grid_map_->gridToWorld(0, grid_map_->getHeight() - 1)[1] - grid_map_->gridToWorld(0, 0)[1]);
		int terrain_x = static_cast<int>(lon_frac * (terrain_data_[0].size() - 1));
		int terrain_y = static_cast<int>(lat_frac * (terrain_data_.size() - 1));
		terrain_x = std::max(0, std::min((int)terrain_data_[0].size() - 1, terrain_x));
		terrain_y = std::max(0, std::min((int)terrain_data_.size() - 1, terrain_y));
		return terrain_data_[terrain_y][terrain_x];
	}
	bool AlgoAStarPrivate::isInside(const algobase::Coord& point, const algobase::Polygon& polygon) const {
		if (polygon.size() < 3) return false;
		int crossings = 0;
		for (size_t i = 0; i < polygon.size(); ++i) {
			const auto& p1 = polygon[i];
			const auto& p2 = polygon[(i + 1) % polygon.size()];
			if (((p1[1] <= point[1] && point[1] < p2[1]) || (p2[1] <= point[1] && point[1] < p1[1])) &&
				(point[0] < (p2[0] - p1[0]) * (point[1] - p1[1]) / (p2[1] - p1[1]) + p1[0])) {
				crossings++;
			}
		}
		return (crossings % 2) != 0;
	}

	double AlgoAStarPrivate::calculateGeographic3DDistance(const algobase::CoordZ& p1, const algobase::CoordZ& p2) const
	{
		return 0.0;
	}

} // namespace algo
